import time
from copy import deepcopy
from get_sample import get_exp_sample, get_template_sample
from Seacher.DocSearcher import DocSearcher
from llm import llm
from kernel import CodeKernel, run_code
import re
import json

doc_searcher = DocSearcher('./data/数据字典.xlsx')
orm_sample = get_exp_sample()
template_sample = get_template_sample()

system_prompt = '''
### 模拟开发工程师查找数据的交互流程

#### 角色介绍
- **用户**：提出问题或请求。
- **助理（您）**：模拟专业金融行业数据开发工程师，负责理解问题、查找文档、编写查询代码并给出答案。
- **IPython**：执行Python代码并返回结果。

#### 行为模式

1. **理解问题**
   - 接收到用户的问题后，基于上下文和之前的对话，将问题标准化并明确化。确保对用户的需求有准确的理解。
   - 禁止向用户提问。

2. **文档查找手册**
   - 使用`doc_searcher`类的`search_by_question`方法，根据用户的原始问题搜索相关的数据库表格和字段，将会展示一个数据库最匹配的top n 个字段。
     - 代码示例：
     ```python
     doc_searcher.search_by_question("市场-用户的原始问题")
     ```
   - 根据搜索结果，选择最合适的表格和字段。
   - 请一定基于IPython返回的文档进行选择，不要遗忘筛选条件，例如日期等。

   - 如果初次搜索结果不充分，使用`next_page` 和 `show_table`方法继续查找更多相关信息。其中 `next_page` 可以查看更多表格信息。`show_table` 可以查看一个表的所有字段

     ```python
     doc_searcher.next_page() # 没有在搜索结果内找到合适的表，你可以调用此函数查看下一页
     ```
     ```python
     doc_searcher.show_table("TableName") # 找到了对应表但是字段不全时可以查看该表
     ```
     ```python
     doc_searcher.show_all_tables() # 如果查询到最后一页依然没有合适的表，你应该查询所有的表，然后再去调用 show_table 方法
     ```
     总结说明：
    1. 使用 `search_by_question` 方法进行初步搜索。
    2. 如果初次搜索结果不充分，使用 `next_page` 和 `show_table` 方法继续查找更多相关信息。
    3. 迫不得已时，使用 `show_all_tables` 方法查看所有表格，然后调用 show_table 方法查看表格内容。


3. **编写查询代码**
   - 使用`sqlalchemy`构建SQL查询语句。仅生成SQL，不要执行。
   - 对于复杂的查询，可以通过子查询来逐步构建查询逻辑。
   - 确保仔细阅读完成数据库文档后开始编写，Table和Column必须使用表内有的。
   - 代码示例：
     ```python
     query = select(Table.Column1, Table.Column2).where(Table.Condition == 'value')
     print("查询xxx开始")
     df_result = exec_orm_sql(query.compile(engine, compile_kwargs={"literal_binds": True}))
     print("查询xxx结束")
     ```
    - 备注:所有的表已经导入，并且类名与数据库文档的表名一致，对象类型是Table，查询直接使用`Table.Column` 注意大小写一致
    - 数据库所有的日期均为字符串类型，请转为日期类型，然后进行查询，
    - 例如过滤交易日:`func.date(CS_HKStockPerformance.TradingDay) == '2020-10-27'`。
    - 过滤年份：`func.year(cast(LC_TransferPlan.InfoPublDate, Date)) == 2021`
    - exec_orm_sql 会自动打印结果，如果你需要其他信息，记得使用print打印

4. **验证查询结果**
   - 执行查询后，检查返回的结果是否符合用户的预期。
   - 如果结果符合要求，以“【答案:xxxx】”的格式给出最终答案，请按照原始内容输出，除非用户有要求，否则不要改变查询答案的格式和内容。
   - 如果结果不符合要求，分析原因并采取相应的措施，如调整查询条件或进一步查找信息。
   - 输出的格式要满足用户的要求，比如日期，小数点，数字不要加分隔符。
   - 关注结果条数是否符合查询，如有些查询应该为1条，有些查询应该为多条。
   - 回答用户参考运行结果，需要重新整理，根据用户的要求返回实体，如涉及公司，则返回`公司名(简称)`的形式，不要返回innercode

5. **注意事项**
    - 你需要重点关注一些选择字段，例如 `IfPerformed`,`IfMerged`,`xxType`等,尤其是有注释的字段，根据题意选择正确的过滤方法。
    - 对于不确定的过滤字段，如需统计，你需要查询并在pandas里进行两次统计，分别回答。
    - 查询XXXName的时候必须查询XXXCode
6. **注意思考**
    - 如果用户需要你仔细思考问题，你会使用 `<thought>你的思考过程</thought>` 进行思考


## SQLAlchemy 编写准则
1. 复杂的查询尽量使用子查询，查询错误后使用分步执行，此时使用 `in_` 代替表连接
2. 按照步骤编写，在每个步骤前，加上注释
3. 打印结果，每个步骤后都要使用print打印结果，`exec_orm_sql` 除外,输出结果只输出前20条，除非确定此结果为答案可全量输出。中间过程禁止打印。
4. 禁止假设，你给出的代码必须不含假设数据，不能空白
5. 可执行，你的代码必须是可以执行的代码，无需用户再填入内容
6. 仔细阅读表文档，谨慎使用 sum count 等函数
## exec_orm_sql 函数说明
 - 执行SQL并返回结果，结果格式为 pandas的DataFrame类
 - 自动打印结果，至多20条
 - 自动将innercode等转化为`公司名(简称)`的形式，

## 金融知识驱动

当用户的问题有金融相关知识背景时，思考硬包含金融名词，如：
概念，行业，价格，月涨幅等，对这些你需要详细解释，比如：
月涨幅：每月最后一个交易日可以计算月涨幅，计算方法为月末价格比月初价格。

#### 精细化解题示例，以下演示了精细化思考解决问题的方法
-精细化示例-
[template]
-----
请注意，以上流程是多轮对话的示例，每一轮对话的输出都依赖于之前的对话内容。请根据用户的请求和系统的反馈逐步进行操作。
'''


class ORMAgent:
    def __init__(self):
        # self.sample = sample
        self.dynamics_prompt = []
        self.current_question = []
        self.history = []
        self.messages = []
        self.save_dict = {}
        self.doc_dict = {}

    def start_code_kernel(self):
        self.ckernel = CodeKernel()
        self.ckernel.start()

        while True:
            time.sleep(1)
            tr = run_code('''print("test output")''', self.ckernel, True)
            # print(f'开始测试输出：```{tr}```')
            if tr == 'test output':
                break
        run_code('''
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "last_expr"
# 简化异常显示
%xmode Minimal
''', self.ckernel, True)
        run_code('''
import warnings
warnings.filterwarnings("ignore")
import pandas as pd
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 20)
pd.set_option('display.float_format', lambda x: str(x))
from Models import *
from sqlalchemy import *
from utils.utils import exec_orm_sql
engine = create_engine('mysql+pymysql://')
''', self.ckernel, True)

    def make_dynamics_prompt(self, question):
        exps = orm_sample.get_sample(question, 5)  # .reverse()
        temp = template_sample.get_sample(question, 1)[0]
        dynamics_prompt = [
            {'role': 'system',
             'content': system_prompt.replace('[template]', temp)
             }
        ]
        exps = '\n---\n'.join(exps)

        dynamics_prompt.extend(
            [{"role": 'user', "content": f"---doc---{exps}\n---\n以上是一些ORM编写示例，你的代码参考此风格，但是对于新问题你要参考精细化示例重新思考，并严格按照`行为模式`执行。"},
             {"role": "assistant", "content": "好的我已知晓以上是示例，我会再次独立思考。"
                                              "我以后会按照精细化示例的风格进行分析。并按照`行为模式`执行。"}])

        self.dynamics_prompt = dynamics_prompt

    def add_question(self, question):
        self.current_question = []
        doc_txt = doc_searcher.search_by_question(question)
        self.doc_dict = {}
        for d in doc_txt.split('\n---\n'):
            k = d.split('\n')[0].split(':')[1]
            self.doc_dict[k] = d
        self.current_question.extend([{'role': 'user',
                                       'content': '请解决新问题：\n' + question + '\n注意`思考`，注意`SQLAlchemy 编写准则`'},
                                      {'role': 'assistant', 'content': f'''<thought>基于文档查找手册，首先我要基于用户问题搜索文档</thought>
```python
doc_searcher.search_by_question("{question}")
```
请执行搜索并返回结果。
'''}, {"role": "user", "content": 'IPython:' + doc_txt + f'''
以上是文档，请分析文档以及用户的问题，选择正确的表格及其字段编写代码回答问题。注意表格和字段说明，如果问题比较复杂可以分多步解决，灵活使用ORM和Python。
步骤说明：
1. 举查询方法：
    - 列举两种可能的查询方法。优先方法：xxx 备用方法：xxx
2. 实现优先方法：
    - 编写代码实现优先方法。
    - 如果遇到困难则更换为备用方法。
3. 单表查询：
    - 主要关注单表查询。
    - 对于过滤条件，分析其必要性。
    - 对于不确定的过滤条件，选择相关过滤字段后，在pandas中进行进一步过滤。
4. 时间和筛选字段：
    - 特别注意时间和筛选字段的分析。
    - 对于历史数据（如股票数据），考虑交易日选择。
    - 对于其他查询，考虑最新发布的数据。
    - 公司状态如网址地址信息无需考虑。
5. 统计类信息：
    - 考虑筛选条件是否适用。
    - 如不确定，可以都统计，但在pandas中进行过滤。
6. 查询字段：
    - 查询xxxName字段时，如果有xxxCode字段，需同时查询xxxCode，例如查询SubclassName时顺带查询SubclassCode。
    - 思考需要查询的字段，并列举。
    - 对于日期过滤，默认为 EndDate，除非有其他合适的 Date
    - 对于股票计算，如涨跌幅，月年最高等，使用表内字段统计的，不要自己统计。 
7. 编码简化：
    - 尽量简化编码，大部分问题单表查询加上一些Python处理即可解决。
    - `exec_orm_sql`返回的是`DataFrame`，判断使用 `not result.empty`。
    - 所有models, sqlalchemy, pandas,engine 已经导入，无需重复import。请分析后编写代码，我将帮你执行代码。
请开始编码：
'''}])

    def solve(self, question_info):
        tid = question_info['id']
        question = question_info['question']
        self.save_dict['final_code'] = []
        for try_num in range(5):
            res = llm(self.dynamics_prompt + self.history + self.current_question)
            self.current_question.append({'role': 'assistant', 'content': res})
            if ('答案' in res or '任务完成' in res) and '```python' not in res:
                self.save_dict['messages'] = deepcopy(self.current_question)
                self.save_dict['tid'] = tid
                self.save_dict['answer'] = res
                cache_li = self.current_question[:2] + self.current_question[-4:]
                cache_li[2]['content'] = '\n'.join(
                    [self.doc_dict[i] for i in self.doc_dict.keys()
                     if i in '\n'.join(self.save_dict['final_code'])])

                cache_li[-2]['content'] = '...' + cache_li[-2]['content'][-700:]
                self.history.extend(cache_li)

                return self.current_question[-2]['content'].split('请判断输出结果能否回答问题')[0] + '\n' + res

            code_list = re.findall(r'```python(.*?)```', res, re.DOTALL)
            run_log = []
            for code in code_list:
                # print(code)
                if 'doc_searcher.choose' in code:
                    code = re.search("doc_searcher.*$", code, re.DOTALL).group()
                    cache_output = []
                    self.save_dict['choose'] = code
                    try:
                        exec(f'r11 = {code}\n\ncache_output.append(r11)')
                    except Exception as e:
                        cache_output.append(str(e))
                    run_log.extend(cache_output)

                elif 'doc_searcher' in code:
                    code = re.search('doc_searcher.*', code).group()
                    cache_output = []
                    try:
                        exec(f'r11 = {code}\n\ncache_output.append(r11)')
                    except  Exception as e:
                        cache_output.append(str(e))
                    run_log.extend(cache_output)
                else:
                    # print('开始执行代码')
                    for _ in range(2):
                        r = run_code(code, self.ckernel, iscode=True)
                        self.save_dict['final_code'].append(code)
                        if r:
                            # print(f'代码执行完毕,结果为：```{r}```')
                            break
                        else:
                            self.ckernel.shutdown()
                            self.start_code_kernel()
                    if 'Empty DataFrame' in r:
                        r += f'''检测到查询为空，请判断空值是否合理，如果合理请回答问题:{question}，如果不合理请重新构造查询，
重新构造查询请使用以下方法：
1. 检查日期是否转换,字符串查询使用 like
2. 减少条件逐步筛选，并分步打印，只留下一个最稳妥的筛选条件，然后在pandas里进行下一步筛选
3. 查询少量数据样例，通过pandas进一步筛选
4. 更换表，更换Date等过滤条件。
5. 考虑此问题数量为0
注意`思考`和`SQLAlchemy 编写准则`'''
                    elif r == '':
                        r += '请使用print打印信息'
                    run_log.append(r)

            return_log = '\n'.join(run_log)
            if not re.search("Empty DataFrame|Tracebcak|Error", return_log, re.DOTALL | re.IGNORECASE):
                return return_log

            if run_log == []:
                run_log = ['错误:代码运行失败或者超时，请优化ORM，尽量使用表外连表分步骤进行。']

            if try_num > 3:
                self.current_question.append({'role': 'user', 'content': 'IPython:' + '\n'.join(run_log) +
                                                                    f'\n请基于所有历史记录尽可能回答问题：`{question}`的内容,如果无法回答任何一点，请回答`【答案：任务失败,请将任务简化】`，否则回答`【答案：可以回答的部分】`'
                                         })
            else:

                self.current_question.append({'role': 'user', 'content': 'IPython:\n---日志---' + '\n'.join(run_log) + f"""
---
请判断输出结果能否回答问题`{question}`。
如果符合要求，请按照`任务完成\n【答案:xxx】`的格式重新整理。注意答案应从日志中提取的原始内容，不要增加或减少任何符号（如数字千分位）。
忽略格式化输出的错误，只要信息足够回答问题即可。注意答案格式以及用户需求的输出格式。
答案应该只有查询到的`值`的部分。对于公司名称，如无明确要求，应该回答`全称(简称)`；如果有要求，请仅回答指定字段。保留小数和日期格式也要注意。
其中日志里的公司名称格式为`全称(简称)`，请仔细判断输出的内容，不要只相信日志答案部分。
如果存在错误或不符合常识，请回答`任务继续`并重新修改代码。修复代码前先在<thought></thought>标签内思考错误，然后将修复后的代码写在```python ```内。
"""
                                             })

    def run(self, question_info):
        self.make_dynamics_prompt(question_info['question'])
        self.add_question(question_info['question'])
        return self.solve(question_info)
